from fastapi import APIRouter

from core.schemas import ResponseModel, MyAiInput
from csv_agent.agent_retrieval import get_retriever, get_retrieve
from my_ai.myai import chain

router = APIRouter(
    prefix="/my_ai",
    tags=["my_ai"],
    responses={404: {"description": "Not found"}},
)


@router.post("/invoke", response_model=ResponseModel)
async def my_ai_invoke(user_input: MyAiInput):
    first_response = chain.invoke({"input": user_input})
    return {"data": first_response, "code": "200", "message": "Response created successfully"}


@router.post("/retriever_invoke", response_model=ResponseModel)
async def retriever_invoke(_input: MyAiInput):
    response = get_retriever(_input.user_input)
    return {"data": response, "code": "200", "message": "Response created successfully"}


@router.post("/retrieve", response_model=ResponseModel)
async def retrieve(_input: MyAiInput):
    response = get_retrieve(_input.user_input)
    return {"data": response, "code": "200", "message": "Response created successfully"}
